package com.xagent.engine.context.phase;

import com.xagent.engine.context.inter.ITaint;
import com.xagent.engine.context.TaintImpl;
import com.xagent.engine.enums.Phase;
import com.xagent.engine.hook.HookInfo;
import com.xagent.engine.hook.MethodInfo;
import com.xagent.engine.policy.MethodContext;
import com.xagent.engine.utils.HookUtil;
import com.xagent.engine.context.ThreadContext;

import java.util.HashMap;

/**
 * Author: 4ut15m
 * 生成传播链路
 */
public class GeneratorPhase {

    public static void generator(HashMap<String, HashMap<String, HookInfo>> policyMap, MethodContext methodContext, ThreadLocal<ThreadContext> tc) {
        if (policyMap == null || methodContext == null) return;
        /**
         * 当线程上下文为null时，只处理source相关内容，非source hook点触发，不跟踪
         * */
        ThreadContext threadContext = null;
        if (tc.get() == null) {
            threadContext = new ThreadContext();
            tc.set(threadContext);
            threadContext.setEnter(true);
        } else threadContext = tc.get();

        if (threadContext.getPhase() == null && !HookUtil.isSource(policyMap, methodContext.getClassName(), methodContext.getSignature()))
            return;
        MethodInfo methodInfo = HookUtil.getHookInfo(policyMap, methodContext.getClassName(), methodContext.getSignature()).getAMethodInfo(methodContext.getSignature());
        methodContext.setId(methodInfo.getId());
        methodContext.setMethodInfo(methodInfo);
        if (threadContext.getPhase() == null) {
            SourcePhase sourcePhase = new SourcePhase(methodContext);
            sourcePhase.setPhaseStatus(Phase.SOURCE);
            threadContext = new ThreadContext();
            threadContext.setPhase(sourcePhase);
            threadContext.setEnter(true);
            tc.set(threadContext);
            // 污染数据如果是Source或者传播点的R即返回值，那么就需要在fill的时候填充进去
            if ("R".equals(methodInfo.getTarget())) {
            }
            if ("O".equals(methodInfo.getTarget())){
                ITaint taint = new TaintImpl(methodContext.getThisObject());

            }

        }
        if (HookUtil.isPropagate(policyMap, methodContext.getClassName(), methodContext.getSignature())) {
            AbstractPhase abstractPhase = threadContext.getPhase();
            PropagatePhase propagatePhase = new PropagatePhase(methodContext);
            propagatePhase.setParentPhase(abstractPhase);
            propagatePhase.setPhaseStatus(Phase.PROPAGATE);
            threadContext.setPhase(propagatePhase);
        }
        if (HookUtil.isSink(policyMap, methodContext.getClassName(), methodContext.getSignature())) {
            AbstractPhase abstractPhase = threadContext.getPhase();
            SinkPhase sinkPhase = new SinkPhase(methodContext);
            sinkPhase.setParentPhase(abstractPhase);
            sinkPhase.setPhaseStatus(Phase.SINK);
            threadContext.setPhase(sinkPhase);
        }


    }

    public static void fill(HashMap<String, HashMap<String, HookInfo>> policyMap, MethodContext methodContext, ThreadLocal<ThreadContext> tc) {
        if (policyMap == null || methodContext == null) return;
        /**
         * 当线程上下文为null时，只处理source相关内容，非source hook点触发，不跟踪
         * */
        ThreadContext threadContext = tc.get();

        if (threadContext.getPhase() == null && !HookUtil.isSource(policyMap, methodContext.getClassName(), methodContext.getSignature()))
            return;
        if (threadContext.getPhase() != null) {
            AbstractPhase phase = threadContext.getPhase();
            while (phase.parentPhase != null || phase.getClass().getSimpleName().contains("Source")) {
                if (phase.getMethodContext().getIndex() == methodContext.getIndex()) {
                    phase.getMethodContext().setReturnObject(methodContext.getReturnObject());
                    break;
                }


                phase = phase.parentPhase;
            }
            /**
             * test start
             * */
            if (phase.getPhaseStatus() == null || phase.getPhaseStatus() == Phase.PROPAGATE) return;
            String className = methodContext.getClassName();
            String methodName = methodContext.getMethodName();
            Object[] args = methodContext.getParamObjects();
            System.out.println("[XAGENT] current className:" + className);
            System.out.println("[XAGENT] current methodName:" + methodName);
            System.out.println("[XAGENT] current args:");
            for (Object arg : args) {
                if (arg != null && arg.getClass().getName().contains("String")) {
                    if (arg.getClass().toString().contains("[Ljava.lang.String;")) {
                        for (String s : (String[]) arg) {
                            System.out.println(s);
                        }
                    } else System.out.println(arg);
                }
            }
            System.out.println("[XAGENT] return Object:" + methodContext.getReturnObject());


            /**
             * test end
             * */
        }


    }


}
